import ast
import unittest

from 程序核心代码.ast节点.ast操作器.AST操作器 import AST操作器


class TestAST操作器(unittest.TestCase):
    def setUp(self):
        self.代码 = """
class MyClass:
    def method(起始接口):
        pass

def foo(x):
    return x + 1

foo(10)
"""
        self.解析后的AST = ast.parse(self.代码)
        self.操作器 = AST操作器(self.解析后的AST)

    def test_插入节点(self):
        赋值节点 = ast.Assign(
            targets=[ast.Name(id='c', ctx=ast.Store(), lineno=1, col_offset=0)],
            value=ast.Constant(value=100, lineno=1, col_offset=0),
            lineno=1,
            col_offset=0
        )
        self.操作器.插入节点(['body'], 赋值节点, 0)
        self.assertEqual(len(self.解析后的AST.body), 4)

    def test_修改节点(self):
        self.操作器.修改节点(['body', 2, 'value', 'args', 0], ast.Name(id='c', ctx=ast.Load(), lineno=3, col_offset=0))
        调用节点 = self.操作器.获取节点(['body', 2])
        self.assertEqual(调用节点.value.args[0].id, 'c')

    def test_添加函数(self):
        self.操作器.添加函数('bar', ['y'], [ast.Pass()])
        self.assertEqual(len(self.解析后的AST.body), 4)
        self.assertEqual(self.解析后的AST.body[3].name, 'bar')

    def test_删除函数(self):
        self.操作器.删除函数('foo')
        self.assertEqual(len(self.解析后的AST.body), 2)

    def test_替换变量名(self):
        self.操作器.替换变量名('x', 'z')
        函数节点 = self.操作器.获取节点(['body', 1])
        self.assertEqual(函数节点.args.args[0].arg, 'z')

    def test_捕获异常(self):
        self.操作器.捕获异常('foo')
        函数节点 = self.操作器.获取节点(['body', 1])
        self.assertIsInstance(函数节点.body[0], ast.Try)

    def test_添加装饰器(self):
        self.操作器.添加装饰器('foo', 'decorator')
        函数节点 = self.操作器.获取节点(['body', 1])
        self.assertEqual(函数节点.decorator_list[0].id, 'decorator')

    def test_添加类(self):
        self.操作器.添加类('NewClass', [], [ast.Pass()])
        self.assertEqual(len(self.解析后的AST.body), 4)
        self.assertEqual(self.解析后的AST.body[3].name, 'NewClass')

    def test_删除类(self):
        self.操作器.添加类('NewClass', [], [ast.Pass()])
        self.操作器.删除类('NewClass')
        self.assertEqual(len(self.解析后的AST.body), 3)

    def test_添加类方法(self):
        self.操作器.添加类方法('MyClass', 'new_method', ['y'], [ast.Pass()])
        类节点 = self.操作器.获取节点(['body', 0])
        self.assertEqual(len(类节点.body), 2)

    def test_删除类方法(self):
        self.操作器.删除类方法('MyClass', 'method')
        类节点 = self.操作器.获取节点(['body', 0])
        self.assertEqual(len(类节点.body), 0)

    def test_添加属性(self):
        self.操作器.添加属性('MyClass', 'attr', 42)
        类节点 = self.操作器.获取节点(['body', 0])
        self.assertIsInstance(类节点.body[0], ast.Assign)
        self.assertEqual(类节点.body[0].targets[0].id, 'attr')

    def test_删除属性(self):
        self.操作器.添加属性('MyClass', 'attr', 42)
        self.操作器.删除属性('MyClass', 'attr')
        类节点 = self.操作器.获取节点(['body', 0])
        self.assertEqual(len(类节点.body), 1)


if __name__ == '__main__':
    unittest.main()
